#include <opencv4/opencv2/opencv.hpp>
#include <filesystem>
#include "cuSlideWinCropImg.h"

int main(int argc, char *argv[])
{
    using namespace std;
    using namespace cv;
    namespace fs = std::filesystem;

    cudaStream_t stream;
    cudaStreamCreate(&stream);

    string img_path = (argc > 1) ? argv[1] : "/data/dog_2k.png";
    string save_dir = "/culn/crop_cpp";
    int win_h = 640, win_w = 640;
    double ratio_h = 0.6, ratio_w = 0.6;
    int num_h = 0, num_w = 0;

    if (!fs::exists(save_dir))
        fs::create_directories(save_dir);

    Mat img_mat = imread(img_path);
    int img_h = img_mat.rows, img_w = img_mat.cols, img_c = img_mat.channels();

    calcNumCrops(img_h, img_w, win_h, win_w, ratio_h, ratio_w, num_h, num_w);

    uint8_t* src_gpu = nullptr, * dst_gpu = nullptr;
    cudaMalloc(&src_gpu, img_h * img_w * img_c * sizeof(uint8_t));
    cudaMalloc(&dst_gpu, num_h * num_w * win_h * win_w * img_c * sizeof(uint8_t));
    cudaMemcpy(src_gpu, img_mat.data, img_h * img_w * img_c * sizeof(uint8_t), cudaMemcpyHostToDevice);

    // calc
    cuSlideWinCropImg_launch(src_gpu, dst_gpu, img_h, img_w, win_h, win_w, ratio_h, ratio_w, num_h, num_w, stream);

    // to host
    uint8_t* dst_cpu = (uint8_t*)malloc(num_h * num_w * win_h * win_w * img_c * sizeof(uint8_t));
    cudaMemcpy(dst_cpu, dst_gpu, num_h * num_w * win_h * win_w * img_c * sizeof(uint8_t), cudaMemcpyDeviceToHost);

    // save
    // Mat img = Mat(num_w * num_h * win_h, win_w, CV_8UC3, dst_cpu);
    // imwrite("/culn/crop_cpp/all.png", img);
    for (int i = 0; i < num_h; ++i)
    {
        for (int j = 0; j < num_w; ++j)
        {
            Mat img_crop = Mat(win_h, win_w, CV_8UC3, dst_cpu + ((i * num_w + j) * (win_h * win_w * img_c)));

            string save_name = to_string(i) + "_" + to_string(j) + ".png";
            string save_path = save_dir + "/" + save_name;
            imwrite(save_path, img_crop);
        }
    }

    cudaFree(src_gpu);
    cudaFree(dst_gpu);
    free(dst_cpu);
    cudaStreamDestroy(stream);
    cout << "finish" << endl;
    return 0;
}
